import numpy as np

def matrix_factorization(R, P, Q, K, steps=5000, alpha=0.01, beta=0.02):
    Q = Q.T
    for step in range(steps):
        for i in range(len(R)):
            for j in range(len(R[i])):
                if R[i][j] > 0:
                    eij = R[i][j] - np.dot(P[i,:], Q[:,j])
                    for k in range(K):
                        P[i][k] = P[i][k] + alpha * (2 * eij * Q[k][j] - beta * P[i][k])
                        Q[k][j] = Q[k][j] + alpha * (2 * eij * P[i][k] - beta * Q[k][j])

        # 计算损失函数的均方根误差（RMSE）
        eR = np.dot(P, Q)
        e = 0
        for i in range(len(R)):
            for j in range(len(R[i])):
                if R[i][j] > 0:
                    e = e + pow(R[i][j] - np.dot(P[i,:], Q[:,j]), 2)
                    for k in range(K):
                        e = e + (beta / 2) * (pow(P[i][k], 2) + pow(Q[k][j], 2))
        rmse = np.sqrt(e / (len(R) * len(R[0])))

        # 判断是否收敛，如果均方根误差小于某个阈值，可以结束迭代
        if rmse < 0.001:
            break

    return P, Q.T

# 创建示例评分矩阵
R = np.array([
    [5, 3, 0, 1],
    [4, 0, 0, 1],
    [1, 1, 0, 5],
    [1, 0, 0, 4],
    [0, 1, 5, 4],
])

# 设置潜在因子数量和迭代步数
K = 2
steps = 5000

# 初始化用户特征矩阵P和物品特征矩阵Q
N = len(R)
M = len(R[0])
P = np.random.randn(N, K)
Q = np.random.randn(M, K)

# 执行矩阵分解
nP, nQ = matrix_factorization(R, P, Q, K, steps=steps)

print("原始评分矩阵R：")
print(R)
print("重构评分矩阵nP*nQ.T：")
print(np.dot(nP, nQ.T))